# -*- coding: utf-8 -*-
"""
Created on Sun Feb  7 16:36:56 2021

@author: wyx
"""


import math
import torch
import torch.nn as nn
import torch.optim as optim
from collections import namedtuple
import random

random.seed(1)

class DQN(nn.Module):
    def __init__(self, n_actions, n_heads):
        super(DQN, self).__init__()
        self.n_actions = n_actions
        self.n_heads = n_heads
        self.fc = nn.Sequential(
            nn.Linear(3+3*n_actions, n_actions, True),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(n_actions, n_actions, True))
        self.actions = nn.Sequential(
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(n_actions, n_actions*n_heads))
        self.state = nn.Sequential(
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(n_actions, n_heads))
    
    def forward(self, x):
        x = self.fc(x)
        x1 = self.actions(x)
        x2 = self.state(x)
        x3 = x2 - x2.mean(dim=1, keepdims=True)
        x4 = x3.reshape(-1, 1).expand(-1, self.n_actions).reshape(-1, self.n_actions*self.n_heads)
        res = x1 + x4
        return res

Transition = namedtuple('Transition', ('state', 'index', 'action', 'next_state', 'reward', 'mask', 'done'))

class ReplayMemory():
    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []
        self.position = 0
        self.leaves = [[0, i] for i in range(self.capacity)]
        self.sumTree = self.initTree()
        self.deep = int(math.log2(capacity))+1
    
    def initTree(self):
        layer = self.leaves
        while len(layer) != 1:
            nodes = []
            for i in range(len(layer) // 2):
                left = layer[2 * i]
                right = layer[2 * i + 1]
                nodes.append([left[0]+right[0], left, right])
            layer = nodes
        return layer[0]
    
    def push(self, weight, *args):
        """
        保存一次交互
        """
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        self.memory[self.position] = Transition(*args)
        index = self.position
        def upgrade(weight, Tree, index, half):
            if len(Tree) == 2:
                Tree[0] = weight
            else:
                a, b = divmod(index, half)
                upgrade(weight, Tree[a+1], b, half//2)
                Tree[0] = Tree[1][0] + Tree[2][0]
        upgrade(weight, self.sumTree, index, self.capacity//2)
        self.position = (self.position + 1) % self.capacity
        
    def sample(self, batch_size):
        total = self.sumTree[0]
        length = total / batch_size
        while True:
            rands = list(map(lambda x:random.uniform(x*length, (x+1)*length), range(batch_size)))
            def get_index(rand, Tree):
                if len(Tree) == 2:
                    return Tree[1]
                else:
                    if rand < Tree[1][0]:
                        return get_index(rand, Tree[1])
                    else:
                        return get_index(rand-Tree[1][0], Tree[2])
            indexes = list(map(lambda x: get_index(x, self.sumTree), rands))
            if max(indexes) < len(self.memory):
                break
        return list(map(lambda x: self.memory[x], indexes))
        # return random.sample(self.memory, batch_size)
    
    def __len__(self):
        return len(self.memory)
        
    
if __name__ == '__main__':
    BATCH_SIZE = 128
    GAMMA = 0.999
    EPS_START = 0.9
    EPS_END = 0.05
    EPS_DECAY = 200
    
    model = DQN(10, 12)
    optimizer = optim.RMSprop(model.parameters())
    memory = ReplayMemory(32768)
    model.eval()
    output = model(torch.ones((32, 33)))